import sys
import numpy as np
import matplotlib.pyplot as plt

# Make sure that caffe is on the python path:
caffe_root = '/home/cgf/caffeLabUsing/'  # this file is expected to be in {caffe_root}/examples
import sys
sys.path.insert(0, caffe_root + 'python')

import caffe

def reconginze(modelName, imgName):
	plt.rcParams['figure.figsize'] = (10, 10)
	plt.rcParams['image.interpolation'] = 'nearest'
	plt.rcParams['image.cmap'] = 'gray'

	import os
	if not os.path.isfile(caffe_root + 'models/' + modelName + '/caffe_train.caffemodel'):
		print("Downloading pre-trained CaffeNet model...")
		sys.exit()
	
	caffe.set_mode_cpu()
	net = caffe.Net(caffe_root + 'models/' + modelName + '/deploy.prototxt', caffe_root + 'models/' + modelName + '/caffe_train.caffemodel', caffe.TEST)

	# input preprocessing: 'data' is the name of the input blob == net.inputs[0]
	transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
	transformer.set_transpose('data', (2,0,1))
	transformer.set_mean('data', np.load(caffe_root + 'python/caffe/' + modelName + '/' + modelName + '_mean.npy').mean(1).mean(1)) # mean pixel
	transformer.set_raw_scale('data', 255)  # the reference model operates on images in [0,255] range instead of [0,1]
	transformer.set_channel_swap('data', (2,1,0))  # the reference model has channels in BGR order instead of RGB

	# set net to batch size of 50
	net.blobs['data'].reshape(50,3,227,227)

	net.blobs['data'].data[...] = transformer.preprocess('data', caffe.io.load_image(imgName))
	out = net.forward()
	result="Predicted class is #{}.".format(out['prob'][0].argmax())
	print(result)

	#plt.imshow(transformer.deprocess('data', net.blobs['data'].data[0]))
	#plt.show()

	# load labels
	model_labels_filename = caffe_root + 'data/' + modelName +'/synset_words.txt'
	try:
		labels = np.loadtxt(model_labels_filename, str, delimiter='\t')
	except:
		labels = np.loadtxt(model_labels_filename, str, delimiter='\t')

	# sort top k predictions from softmax output
	top_k = net.blobs['prob'].data[0].flatten().argsort()[-1:-6:-1]
	print labels[top_k]
	return labels[top_k]

if __name__=='__main__':
	if len(sys.argv) != 3:
		print "Usage: python caffeRecongination.py modelName{such as oxfordCatsAndDogs} imgName{cat.jpg}"
		sys.exit()

	modelName=sys.argv[1]
	imgName=sys.argv[2]
	reconginze(modelName, imgName)
	#imgName="/home/cgf/Abyssinian_1.jpg"
